import json

import polars as pl
import os
import glob
from concurrent.futures import ThreadPoolExecutor


# 获取指定文件夹中所有后缀为xlsx的文件的绝对路径
def get_xlsx_files(directory):
    # 获取指定文件夹中所有后缀为xlsx的文件的绝对路径
    xlsx_files = glob.glob(os.path.join(directory, '*.xlsx'))
    # 将相对路径转换为绝对路径
    xlsx_files = [os.path.abspath(file) for file in xlsx_files]
    return xlsx_files

# 获取子库分类表
def get_sub():
    df = (pl.read_excel('D:/保税坏件货量计算/子库分类.xlsx', sheet_name="Sheet1",
                        read_options={'use_columns': ['子库类型', '子库编码']},  # 只读取特定列的数据
                        engine="calamine")
          .unique(('子库类型', '子库编码'), keep='first', maintain_order=True)
          )
    return df

def read_xlsx_file():
    directory = r'D:\保税坏件货量计算\ERP交易'
    xlsx_files = get_xlsx_files(directory)

    def process_file(file):
        df1 = pl.read_excel(file, sheet_name="MinvAnalysisTransaction",
                            read_options={
                                'use_columns': ["物料编码", "子库编码", "货位编码", "对方子库编码", "对方货位编码",
                                                "交易日期", "交易来源单据号", "交易数量", "来源", "交易类型"],
                                'header_row': 1
                            },
                            schema_overrides={"物料编码": pl.String, "子库编码": pl.String, "货位编码": pl.String,
                                              "对方子库编码": pl.String, "对方货位编码": pl.String,
                                              "交易来源单据号": pl.String, "来源": pl.String, "交易类型": pl.String,
                                              "交易数量": pl.Int64, "交易日期": pl.String},
                            engine="calamine"
                            ).with_columns(
            pl.col("交易日期").str.split_exact(" ", 1).struct.field("field_0").alias("交易日期")
        )
        return df1

    with ThreadPoolExecutor() as executor:
        df_list = list(executor.map(process_file, xlsx_files))

    # 合并所有数据帧
    df = pl.concat(df_list, how="vertical")

    return df

def WAT(df_sub, df):
    wat_sub = df_sub.filter(pl.col("子库类型") == "WAT")
    wat_df = (
        df.join(wat_sub, on="子库编码", how="inner")
        .filter((~pl.col('来源').str.contains('调账')) & (~pl.col('来源').str.contains('返账')))
        .filter(~pl.col('货位编码').str.contains('MJ'))
    ).select(
        '物料编码', '子库编码', '货位编码', '对方子库编码', '对方货位编码', '交易日期', '交易来源单据号',
        '交易数量'
    )
    # 月货量
    # HT 开头过滤
    ht_df = wat_df.filter(pl.col('交易来源单据号').str.starts_with('HT'))
    month_df = ht_df.select(
        pl.col('子库编码', '交易数量')
    ).with_columns(
        pl.lit(1).alias('项数')
    )
    # <交易来源单据号>中"HT"开头不需要做透视表，正数算(验收)，负数算(验收入库)

    ht_df = ht_df.with_columns([
        pl.when(pl.col('交易数量') > 0).then(pl.lit('验收')).otherwise(pl.lit('验收入库')).alias('计费类型'),
        pl.col('交易数量').abs().alias('交易数量'),
        pl.lit(1).alias('项数')
    ]).group_by(['交易日期', '计费类型']).agg([
        pl.col('交易数量').sum().alias('交易数量'),
        pl.col('项数').sum().alias('项数')
    ])
    # HQ0D2 开头过滤
    HQ0D2_df = wat_df.filter(pl.col('货位编码').str.starts_with('HQ0D2')).with_columns([
        pl.when(pl.col('交易数量') > 0).then(pl.lit('验收')).otherwise(pl.lit('验收入库')).alias('计费类型'),
        pl.col('交易数量').abs().alias('交易数量')
    ]).group_by([
        '物料编码', '子库编码', '货位编码', '对方子库编码', '对方货位编码', '交易日期', '交易来源单据号', '计费类型'
    ]).agg([
        pl.col('交易数量').sum().alias('交易数量')
    ]).with_columns(
        pl.lit(1).alias('项数')
    )
    month_df=pl.concat(
        [
            month_df,
            HQ0D2_df.select(
                pl.col('子库编码', '交易数量', '项数')
            )
        ], how="vertical"
    )
    HQ0D2_df=HQ0D2_df.select('交易日期', '计费类型', '交易数量','项数')
    # HJ 包含过滤
    hj_df = wat_df.filter(pl.col('货位编码').str.contains('HJ')).with_columns([
        pl.when(pl.col('交易数量') > 0).then(pl.lit('好转坏验收')).otherwise(pl.lit('好转坏入库')).alias('计费类型'),
        pl.col('交易数量').abs().alias('交易数量')
    ]).group_by([
        '物料编码', '子库编码', '货位编码', '对方子库编码', '对方货位编码', '交易日期', '交易来源单据号', '计费类型'
    ]).agg([
        pl.col('交易数量').sum().alias('交易数量')
    ]).with_columns(
        pl.lit(1).alias('项数')
    )
    month_df = pl.concat(
        [
            month_df,
            hj_df.select(
                pl.col('子库编码', '交易数量', '项数')
            )
        ], how="vertical"
    )
    hj_df = hj_df.select('交易日期', '计费类型', '交易数量', '项数')

    # 合并数据
    dfs = pl.concat([ht_df, HQ0D2_df, hj_df], how="vertical")
    return dfs, month_df

def SHP(df_sub, df):
    SHP_sub = df_sub.filter(pl.col("子库类型") == "SHP")
    shp_df = (
        df.join(SHP_sub, on="子库编码", how="inner")
        .filter(
            (~pl.col('对方货位编码').str.starts_with('HQ0*G')) &
            (~pl.col('对方货位编码').str.starts_with('HQ0G2')) &
            (pl.col('对方货位编码') != '')
        )
    ).select(
        '子库编码', '交易数量', '交易日期', '交易来源单据号', '货位编码',
        '物料编码', '对方子库编码', '对方货位编码'
    )

    # 月货量
    # HT 开头过滤
    ht_df = shp_df.filter(pl.col('交易来源单据号').str.starts_with('HT'))
    month_df = ht_df.select('子库编码', '交易数量').with_columns(
        pl.lit(1).alias('项数')
    )

    # HT 开头过滤
    ht_df = ht_df.with_columns([
        pl.when(pl.col('交易数量') > 0).then(pl.lit('出库')).otherwise(pl.lit('入库')).alias('计费类型'),
        pl.col('交易数量').abs().alias('交易数量'),
        pl.lit(1).alias('项数')
    ]).group_by(['交易日期', '计费类型']).agg([
        pl.col('交易数量').sum().alias('交易数量'),
        pl.col('项数').sum().alias('项数')
    ])

    # RTV 开头过滤
    RTV_df = shp_df.filter(pl.col('交易来源单据号').str.starts_with('RTV')).with_columns([
        pl.when(pl.col('交易数量') > 0).then(pl.lit('出库')).otherwise(pl.lit('入库')).alias('计费类型'),
        pl.col('交易数量').abs().alias('交易数量')
    ]).group_by([
        '物料编码', '子库编码', '货位编码', '对方子库编码', '对方货位编码', '交易日期', '交易来源单据号', '计费类型'
    ]).agg([
        pl.col('交易数量').sum().alias('交易数量')
    ]).with_columns(
        pl.lit(1).alias('项数')
    )
    month_df = pl.concat(
        [
            month_df,
            RTV_df.select(
                pl.col('子库编码', '交易数量', '项数')
            )
        ], how="vertical"
    )
    RTV_df = RTV_df.select('交易日期', '计费类型', '交易数量', '项数')

    # 合并数据
    dfs = pl.concat([ht_df, RTV_df], how="vertical")
    return dfs, month_df

def BF(df_sub, df):
    BF_sub = df_sub.filter(pl.col("子库类型") == "BF")
    bf_df = (
        df.join(BF_sub, on="子库编码", how="inner")
        .filter(
            (pl.col('交易类型') == 'SP_S20') |
            (pl.col('交易类型') == 'SP_OB005')
        )
    ).select(
        '物料编码', '子库编码', '货位编码', '交易日期', '交易来源单据号',
        '交易数量'
    )

    # 报废出库过滤
    bf_df = bf_df.with_columns([
        pl.lit('报废出库').alias('计费类型'),
        pl.col('交易数量').abs().alias('交易数量')
    ]).group_by([
        '物料编码', '子库编码', '货位编码', '交易日期', '交易来源单据号', '计费类型'
    ]).agg([
        pl.col('交易数量').sum().alias('交易数量')
    ]).with_columns(
        pl.lit(1).alias('项数')
    )
    # 月货量
    month_df=bf_df.select(
        pl.col('子库编码', '交易数量', '项数')
    )
    bf_df=bf_df.select('交易日期', '计费类型', '交易数量','项数')

    return bf_df, month_df

def process_functions(df_sub, df):
    def handle_function(func, df_sub, df):
        try:
            return func(df_sub, df)
        except Exception as e:
            return f'错误：处理{func.__name__}数据时发生错误：{e}', None

    functions = [WAT, SHP, BF]
    with ThreadPoolExecutor() as executor:
        results = list(executor.map(lambda func: handle_function(func, df_sub, df), functions))

    # 如果有错误信息，直接返回
    for result in results:
        if isinstance(result, str):
            return result

    # 解包结果
    wat_dfs, wat_month_df = results[0]
    shp_dfs, shp_month_df = results[1]
    bf_dfs, bf_month_df = results[2]

    # 合并结果
    dfs = pl.concat([wat_dfs, shp_dfs, bf_dfs], how="vertical")
    month_df = pl.concat([wat_month_df, shp_month_df, bf_month_df], how="vertical")

    return dfs, month_df


def function1():
    try:
        df_sub = get_sub()
        if df_sub.height == 0:
            return "子库分类表为空"

        df = read_xlsx_file()
        if df.height == 0:
            return "ERP交易表为空"

        # 填充空值
        df = df.fill_null('').with_columns(
            pl.col('交易数量').cast(pl.Int64).alias('交易数量')
        )
        df_sub = df_sub.fill_null('')

        # 并行处理函数
        dfs, month_df = process_functions(df_sub, df)
        if isinstance(dfs, str):  # 如果有错误信息
            return dfs

        # 处理汇总和排序
        dfs = (dfs
        .group_by('交易日期', '计费类型')
        .agg(
            pl.col('交易数量').sum().alias('交易数量'),
            pl.col('项数').sum().alias('项数')
        )
        .sort(['交易日期', '计费类型'])
        .with_columns(
            pl.col('交易数量').cast(pl.Int64),
            pl.col('项数').cast(pl.Int64)
        ))

        summary_row = pl.DataFrame({
            "交易日期": [""],
            "计费类型": ["汇总"],
            "交易数量": [dfs["交易数量"].sum()],
            "项数": [dfs["项数"].sum()]
        })
        dfs = pl.concat([summary_row, dfs], how="vertical")

        month_df = (month_df
        .with_columns(pl.col('交易数量').abs().alias('交易数量'))
        .group_by('子库编码')
        .agg(
            pl.col('交易数量').sum().alias('交易数量'),
            pl.col('项数').sum().alias('项数')
        )
        .sort('子库编码')
        .with_columns(
            pl.col('交易数量').cast(pl.Int64),
            pl.col('项数').cast(pl.Int64)
        ))

        m_row = pl.DataFrame({
            "子库编码": ["汇总"],
            "交易数量": [month_df["交易数量"].sum()],
            "项数": [month_df["项数"].sum()]
        })

        month_df = pl.concat([m_row, month_df], how="vertical")
        # 转换为JSON
        rdf = {
            'dfs': json.dumps(dfs.to_dicts(), ensure_ascii=False),
            'month_df': json.dumps(month_df.to_dicts(), ensure_ascii=False)
        }
        return json.dumps(rdf, ensure_ascii=False)

    except Exception as e:
        return f'错误：处理数据失败：{e}'